package main.java.state.key.casedemo;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.shaded.guava18.com.google.common.collect.Lists;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.Collections;

public class ListCaseDemo2 {
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Tuple2<Long, Long>> dateStream = env.fromElements(Tuple2.of(1L, 3L), Tuple2.of(1L, 5L), Tuple2.of(1L, 7L),
                Tuple2.of(2L, 4L), Tuple2.of(2L, 2L), Tuple2.of(2L, 5L));
        dateStream.keyBy(0)
                .flatMap(new ListDemo2())
                .print();
        env.execute("test");
    }


    public static class ListDemo2 extends RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>> {


        private ListState<Tuple2<Long, Long>> listState;

        @Override
        public void open(Configuration parameters) throws Exception {

            ListStateDescriptor<Tuple2<Long, Long>> descriptor = new ListStateDescriptor<>("test", Types.TUPLE(Types.LONG, Types.LONG));
            listState = getRuntimeContext().getListState(descriptor);

        }

        @Override
        public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Double>> out) throws Exception {


            Iterable<Tuple2<Long, Long>> iterable = listState.get();
            if (null == iterable) {
                listState.addAll(Collections.emptyList());
            }
            listState.add(element);
            ArrayList<Tuple2<Long, Long>> allElements = Lists.newArrayList(iterable);
            if (allElements.size() >= 3) {
                Long count = 0L;
                Long sum = 0L;
                for (Tuple2<Long, Long> e : allElements) {
                    count++;
                    sum += e.f1;
                }
                double avg = (double) sum / count;
                out.collect(Tuple2.of(element.f0, avg));
                listState.clear();
            }


        }
    }
}
